import subprocess
import sys
import json
import os
import logging
from dotenv import load_dotenv

load_dotenv()


def setup_logging():
    """设置日志配置"""
    log_level = os.getenv("LOG_LEVEL", "INFO").upper()
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    # 创建控制台日志处理器，输出到标准错误
    console_handler = logging.StreamHandler(sys.stderr)
    console_handler.setFormatter(formatter)
    # 配置日志记录器
    logging.basicConfig(
        level=getattr(logging, log_level),
        handlers=[console_handler],
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    return logging.getLogger(__name__)


logger = setup_logging()


class MCPClient:
    def __init__(self):
        self.process = None
        self.request_id = 1
        self.available_tools = []
        self.available_resources = []
        self.available_templates = []
        self.available_prompts = []
        self._connected = False
        self._shutdown = False

    def connect_to_server(self, server_script_path):
        if not os.path.exists(server_script_path):
            raise FileNotFoundError(f"服务器脚本不存在:{server_script_path}")
        if not server_script_path.endswith(".py"):
            logger.warning(f"服务器脚本可能不是python文件:{server_script_path}")
        self.process = subprocess.Popen(
            [sys.executable, server_script_path],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=0,
        )
        self._load_all_capabilities()
        self._connected = True

    def _initialize_server(self):
        logger.info("发送初始化请求...")
        init_response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-06-18",
                    "capabilities": {},
                    "clientInfo": {
                        "name": "MCP Client",
                        "title": "MCP Client Display Name",
                        "version": "1.0.0",
                    },
                },
            }
        )
        self.request_id += 1
        if "result" not in init_response:
            logger.error(f"初始化失败:{init_response}")
            return False
        logger.info(f"初始化成功，发送initialized通知")
        self._send_request(
            {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}
        )
        logger.info(f"服务器连接成功")
        return True

    def _send_request(self, request):
        # 检查服务器进程是否存在并且尚未退出
        if not self.process or self.process.poll() is not None:
            logger.error("服务器已经断开")
            return {"error": "Server connection lost"}
        request_str = json.dumps(request, ensure_ascii=False)
        try:
            self.process.stdin.write(request_str + "\n")
            self.process.stdin.flush()
        except Exception as e:
            logger.error(f"发送请求失败:{str(e)}")
            return {"error": f"发送请求失败:{str(e)}"}
        if "id" not in request:
            logger.info(f"通知类的请求，不用等待响应")
            return {"result": "通知已经发送"}
        try:
            response_line = self.process.stdout.readline()
            if response_line:
                try:
                    response = json.loads(response_line.strip())
                    return response
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析错误:{str(e)}")
                    return {"error": f"JSON解析错误:{str(e)}"}
            else:
                logger.error("没有收到响应")
                return {"error": "没有响应"}
        except Exception as e:
            logger.error(f"接收响应失败:{str(e)}")
            return {"error": f"接收响应失败:{str(e)}"}

    def load_resources(self):
        logger.info(f"获取资源列表...")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "resources/list",
                "params": {},
            }
        )
        self.request_id += 1
        if "result" in response and "resources" in response["result"]:
            self.available_resources = response["result"]["resources"]
            logger.info(
                f"可用资源:{[resource['name'] for resource in self.available_resources]}"
            )
        else:
            logger.error(f"获取可用资源列表失败:{response}")

    def load_prompts(self):
        logger.info(f"获取提示词列表...")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "prompts/list",
                "params": {},
            }
        )
        self.request_id += 1
        if "result" in response and "prompts" in response["result"]:
            self.available_prompts = response["result"]["prompts"]
            logger.info(
                f"可用提示词:{[prompt['name'] for prompt in self.available_prompts]}"
            )
        else:
            logger.error(f"获取可用提示词列表失败:{response}")

    def load_templates(self):
        logger.info(f"获取模板列表...")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "resources/templates/list",
                "params": {},
            }
        )
        self.request_id += 1
        if "result" in response and "resourceTemplates" in response["result"]:
            self.available_templates = response["result"]["resourceTemplates"]
            logger.info(
                f"可用模板:{[template['name'] for template in self.available_templates]}"
            )
        else:
            logger.error(f"获取可用模板列表失败:{response}")

    def _load_all_capabilities(self):
        logger.info("正在加载服务器的能力")
        if not self._initialize_server():
            return
        # 加载工具列表
        self._load_tools()
        self.load_resources()
        self.load_templates()
        self.load_prompts()
        logger.info("能力加载完成")

    def _load_tools(self):
        logger.info("加载工具列表")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "tools/list",
                "params": {},
            }
        )
        self.request_id += 1
        if "result" in response and "tools" in response["result"]:
            self.available_tools = response["result"]["tools"]
            logger.info(f"可用工具:{[tool['name'] for tool in self.available_tools]}")
        else:
            logger.error(f"获取工具列表失败:{response}")

    def _ping(self):
        response = self._send_request(
            {"jsonrpc": "2.0", "id": self.request_id, "method": "ping", "params": {}}
        )
        self.request_id += 1
        return response

    def _list_tools(self):
        tool_list = []
        for tool in self.available_tools:
            name = tool.get("name")
            description = tool.get("description")
            tool_list.append(f"- {name} :{description}")
        return f"可用工具:\n" + "\n".join(tool_list)

    def call_tool(self, tool_name, arguments):
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": arguments,
                },
            }
        )
        self.request_id += 1
        return response

    def _call_tool(self, tool_name, args):
        arguments = {}
        if tool_name == "get_current_time":
            if args and args[0] in ["iso", "timestamp", "readable"]:
                arguments["format"] = args[0]
        elif tool_name == "calculate":
            if args and args[0]:
                arguments["expression"] = args[0]
            else:
                return "calculate工具必须提供表达式参数 "
        elif tool_name == "echo":
            if args and args[0]:
                arguments["message"] = args[0]
            else:
                return "echo工具必须提供消息参数 "
        logger.info(f"调用工具:{tool_name},参数:{arguments}")
        response = self.call_tool(tool_name, arguments)
        if "result" in response and "content" in response["result"]:
            logger.info(f"{tool_name}调用成功")
            return response["result"]["content"]
        else:
            logger.info(f"{tool_name}调用失败")
            return f"{tool_name}调用失败:{response}"

    def _list_resources(self):
        resources_list = []
        for resource in self.available_resources:
            name = resource.get("name", "Unknown")
            uri = resource.get("uri", "No URI")
            resources_list.append(f" - {name}:{uri}")
        return f"可用的资源:\n" + "\n".join(resources_list)

    def _read_resource(self, uri):
        resource = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "resources/read",
                "params": {"uri": uri},
            }
        )
        self.request_id += 1
        return resource["result"]

    def _list_templates(self):
        templates_list = []
        for template in self.available_templates:
            name = template.get("name", "Unknown")
            uri = template.get("uri", "No URI")
            templates_list.append(f" - {name}:{uri}")
        return f"可用的模板:\n" + "\n".join(templates_list)

    def _read_template(self, uri):
        resource = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "resources/templates/read",
                "params": {"uri": uri},
            }
        )
        self.request_id += 1
        return resource["result"]

    def _list_prompts(self):
        prompts_list = []
        for prompt in self.available_prompts:
            name = prompt.get("name", "Unknown")
            description = prompt.get("description", "No description")
            prompts_list.append(f" - {name}:{description}")
        return f"可用的提示词:\n" + "\n".join(prompts_list)

    def _parse_arguments(self, args):
        arguments = {}
        for arg in args:
            if "=" in arg:
                key, value = arg.split("=", 1)
                arguments[key] = value
        return arguments

    def _read_prompt(self, prompt_name, arguments):
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "prompts/get",
                "params": {"name": prompt_name, "arguments": arguments},
            }
        )
        self.request_id += 1
        return response["result"]

    def mcp_sampling_command(self):
        logger.info("通知服务器开始采样")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "mcp/sampling",
                "params": {},
            }
        )
        # ---------------------------------------------------
        self.request_id += 1
        params = response.get("params")
        method = response.get("method")
        id = response.get("id")
        messages = params.get("messages")
        modelPreferences = params.get("modelPreferences")
        systemPrompt = params.get("systemPrompt")
        maxTokens = params.get("maxTokens")
        logger.info(f"method:{method}")
        logger.info(f"id:{id}")
        logger.info(f"modelPreferences:{modelPreferences}")
        logger.info(f"systemPrompt:{systemPrompt}")
        logger.info(f"maxTokens:{maxTokens}")
        print("\n请选择操作")
        print("y.批准生成内容")
        print("n.拒绝请求")
        try:
            choice = input("请输入选择(y/n)").strip()
        except (EOFError, KeyboardInterrupt):
            logger.info("用户中断了采样请求")
            return
        if choice == "y":
            print("用户已经批准了采样请求")
            choiceModel = "deepseek-chat"
            # 1.根据modelPreferences选择合适的大模型
            # 2.调用大模型生成结果  messages systemPrompt maxTokens
            sample_result = "The capital of France is Paris."
            print(f"这是大模型返回的结果:", sample_result)
            try:
                choice = input("结果是否可用(y/n)").strip()
            except (EOFError, KeyboardInterrupt):
                logger.info("用户中断了采样请求")
                return
            if choice == "y":
                print("用户已经批准了把采样结果发回给MCP服务器的请求")
                # 把大模型生成的结果发回给mcp服务器
                result = {
                    "jsonrpc": "2.0",
                    "id": id,
                    "method": "sampling/result",
                    "result": {
                        "role": "assistant",
                        "content": {
                            "type": "text",
                            "text": sample_result,
                        },
                        "model": choiceModel,
                        "stopReason": "endTurn",
                    },
                }
                self._send_request(result)
            else:
                print(f"用户已经拒绝向服务器返回采样结果")
        else:
            print(f"用户已经拒绝了采样请求")

    def _mcp_elicitation_command(self):
        logger.info("执行演示引导命令")
        response = self._send_request(
            {
                "jsonrpc": "2.0",
                "id": self.request_id,
                "method": "mcp/elicitation",
                "params": {},
            }
        )
        self.request_id += 1
        self._handle_elicitation_request(response.get("id"), response.get("params"))

    def _handle_elicitation_request(
        self,
        request_id,
        params,
    ):
        logger.info("收到服务器的启动请求")
        print("\n服务器请求用户的偏好设置")
        print(f"消息:{params.get('message')}")
        print(f"服务器希望用户填写的信息:{params.get("requestedSchema")}")

    def process_query(self, query):
        # 去掉首尾的空格，并按空格分割 quit help  read resource
        parts = query.strip().split()
        if not parts:
            return "请输入有效的命令"
        command = parts[0].lower()
        try:
            if command == "help":
                return self._show_help()
            elif command == "ping":
                result = self._ping()
                return f"Ping的结果为:{result}"
            elif command == "tools":
                return self._list_tools()
            elif command in ["get_current_time", "calculate", "echo"]:
                return self._call_tool(command, parts[1:])
            elif command == "resources":
                return self._list_resources()
            elif command == "read":
                if len(parts) < 2:
                    return "用法 read <资源ID>"
                uri = parts[1]
                result = self._read_resource(uri)
                return f"资源内容:{json.dumps(result,ensure_ascii=False,indent=2)}"
            elif command == "templates":
                return self._list_templates()
            elif command == "template":
                if len(parts) < 2:
                    return "用法 template <模板ID>"
                uri = parts[1]
                result = self._read_template(uri)
                return f"模板内容:{json.dumps(result,ensure_ascii=False,indent=2)}"
            elif command == "prompts":
                return self._list_prompts()
            elif command == "prompt":
                if len(parts) < 2:
                    return "用法 prompt <提示词名称> [参数]"
                prompt_name = parts[1]
                arguments = self._parse_arguments(parts[2:])
                result = self._read_prompt(prompt_name, arguments)
                return f"提示结果:{json.dumps(result,ensure_ascii=False,indent=2)}"
            elif command == "mcp-sampling":
                return self.mcp_sampling_command()
            elif command == "mcp-elicitation":
                return self._mcp_elicitation_command()
            else:
                return f"未知命令{command}\n请输入help查看帮助"
        except Exception as e:
            logger.error(f"执行命令时出错了:{str(e)}")

    def _show_help(self):
        return """
        === MCP客户端帮助 ===

            基础命令:
            help          - 显示此帮助
            ping          - 测试服务器连接

            工具相关:
            tools         - 列出所有工具
            - get_current_time <format>  获取 当前的时间(format : iso/timestamp/readable)
            - calculate <expression>     计算表达式的结果 
            - echo <message>             回显消息

            资源相关:
            resources  - 列出所有的资源
            read <uri> - 读取资源内容

            url模板相关:
            templates  - 列出所有的模板
            template <uri> - 读取模板内容

            提示词相关:
            prompts  - 列出所有的提示词
            prompt <name> [args] - 使用提示词

            演示功能：
            mcp-sampling - 演示AI内容生成功能
            mcp-elicitation - 演示用户偏好收集功能

            示例:
            help
            ping 
            tools
            resources
            read mcp://server-info
            templates
            template mcp://templates/project-plan
            prompts
            prompt greeting name=Alice time_of_day=morning
            mcp-sampling
            mcp-elicitation
        """

    def chat_loop(self):
        logger.info(f"MCP客户端已经启动")
        print("输入 help 查看所有的命令")
        print("输入 quit 退出")
        while True:
            try:
                query = input("\n命令:").strip()
                if query.lower() == "quit":
                    logger.info("用户想退出客户端")
                    break
                response = self.process_query(query)
                print(f"{response}")
            except Exception as e:
                logger.error(f"处理命令时出错:{str(e)}")

    def cleanup(self):
        self._shutdown = True
        self._connected = False
        try:
            if self.process:
                self.process.terminate()
                try:
                    self.process.wait(timeout=5)
                    logger.info("服务器进程正常终止")
                except subprocess.TimeoutExpired:
                    logger.warning(f"服务器进程未在5秒内终止，强制杀死")
                    self.process.kill()
                    self.process.wait()
                    logger.warning(f"服务器进程已被强行终止")
        except Exception as e:
            logger.error(f"清理资源出错:{str(e)}")
        finally:
            self.process = None


def main():
    if len(sys.argv) < 2:
        logger.error(f"用法: python client.py <服务器脚本路径>")
        print("用法: python client.py <服务器脚本路径>")
        sys.exit(1)
    logger.info(f"启动MCP客户端，连接服务器:{sys.argv[1]}")
    client = MCPClient()

    try:
        client.connect_to_server(sys.argv[1])
        client.chat_loop()
    except Exception as e:
        logger.error(f"客户端运行出错:{str(e)}")
    finally:
        logger.info("清理客户端的资源")
        client.cleanup()


# python client.py server.py
if __name__ == "__main__":
    main()
